from datasets.pair_match_dataset import PairMatchDataset
from tools.QXSLAB_tools import QXSLAB_tools
from tools.sentinel12_tools import sentinel12_tools
from tools.os_tools import os_tools


class Sentinel12PairMDataset(PairMatchDataset):


    def __init__(self, config=None):
        super(Sentinel12PairMDataset, self).__init__(
            config, tools=sentinel12_tools, name="Sentinel12"
        )


class OSPairMDataset(PairMatchDataset):


    def __init__(self, config=None):
        super(OSPairMDataset, self).__init__(
            config, tools=os_tools, name="OS"
        )


class QXSLABPairMDataset(PairMatchDataset):


    def __init__(self, config=None):
        super(QXSLABPairMDataset, self).__init__(
            config, tools=QXSLAB_tools, name="QXSLAB"
        )


# test datasets
if __name__ == "__main__":
    import cv2 as cv
    from dotmap import DotMap
    import numpy as np

    # dset_kwt = "os"
    # list_file = "E:/workspace/SOMatch/json/os_dataset/256_train.json"
    # dset_obj = OSPairMDataset

    # dset_kwt = "sen12"
    # list_file = "E:/workspace/SOMatch/json/sen12_summer_harris/HARRIS_o1s1.json"
    # dset_obj = Sentinel12PairMDataset

    dset_kwt = "QXSLAB"
    # list_file = "E:/workspace/SOMatch/json/QXSLAB_SAROPT/QingDao.json"
    list_file = "/home/rslab/DLMatch/workspace/SOMatch/json/despeckle.json"
    dset_obj = QXSLABPairMDataset

    class config:

        def __init__(self):
            self.single_domain = ""
            self.cache_dir = None
            self.cache_size = 0
            self.base_dir = "/home/rslab/DLMatch/datasets/despeckle/1"
            self.list_file = list_file
            self.augment = {
                "fliplr": 0,
                "flipud": 0,
                "scale": 0,
                "scale_px": (1.0, 1.0),
                "translate": 0,
                "translate_perc": (0.0, 0.0),
                "rotate": 0,
                "rotate_angle": (-5, 5)
            }
            # self.clipper = DotMap({
            #         "name": "center same",
            #         "template_size": 128
            # })

    cfg = config()

    ppd = dset_obj(cfg)

    img0, img1 = ppd[10]
    img0 = img0.numpy()
    img1 = img1.numpy()
    img0 = np.transpose(img0, [1, 2, 0])
    img1 = np.transpose(img1, [1, 2, 0])

    cv.imwrite(f"./tmp/{dset_kwt}_pariM_sar.png", img0*255)
    cv.imwrite(f"./tmp/{dset_kwt}_pariM_opt.png", img1*255)


